Repository navigation
Extend MoEPacked for Bolt ReLU-square routing and BF16 IQ2 precision - #33963
Open
Rui Ren (rui-ren) wants to merge 9 commits into
Open
Rui Ren (rui-ren) wants to merge 9 commits into
Rui Ren (rui-ren) wants to merge 9 commits into
Conversation
Rui Ren (rui-ren)
force-pushed
the
rui-ren/bolt-packed-moe-on-33953-20261107
branch
from
November 7, 2026 05:42
520f375 to
b9b776b
Compare
Rui Ren (rui-ren)
changed the base branch from
tlwu/qwen38-gguf
to
rui-ren/bolt-packed-projections-on-33953-20261107
November 7, 2026 05:43
3 of 4 tasks
Rui Ren (rui-ren)
force-pushed
the
rui-ren/bolt-packed-moe-on-33953-20261107
branch
from
November 7, 2026 06:34
b9b776b to
c818493
Compare
3 of 4 tasks
Rui Ren (rui-ren)
force-pushed
the
rui-ren/bolt-packed-moe-on-33953-20261107
branch
2 times, most recently
from
November 7, 2026 17:24
736d86e to
a9bdbce
Compare
This comment was marked as resolved.
This comment was marked as resolved.
This was referenced Nov 8, 2026
Contributor
Author
|
Standalone Bolt routing is now isolated in draft follow-up #33994 above #33969, extending native stack #33967 without rewriting this PR or the coworker's #32935. The existing integrated route remains compatible and uses the same shared CUDA implementation. The complete fresh production source build passed native tests, logits/cache/state and route comparisons, plus 240 exact same-mode generation/reset/counter checks. The repeated 12k/24k/32k/124k before/after timings and limitations are in #33994; those timings use the frozen extraction build, not a new benchmark of the production refactor. Preserved quantized control is not independent HF, and inherited ordinary-versus-k7 differences remain. |
Contributor
Author
|
The native kernel optimizations are now pushed in #34002 and registered as the Actual #33994-tip baseline versus the new native candidate, real 1028-1663-token The earlier ~92 tok/s private baseline already contained these optimizations: |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The BF16 IQ2 ReLU-square tests leave channel permutations and FC2 cancellation-sensitive accumulation order untested, requiring asymmetric reference cases and larger mixed-sign fixtures for both decode and prefill.
2 open issues
What changed in this PR
Extends the existing CUDA MoEPacked operator for Nemotron Bolt’s ReLU-square experts and routing requirements while retaining default SwiGLU behavior.
Changes:
- Adds preselected FP32 routes and projected sigmoid routing with selection-only correction bias.
- Introduces a BF16 IQ2 kernel preserving contiguous-K accumulation and rounding boundaries.
- Adds routing and precision regression tests and updates operator documentation.
| File | Description |
|---|---|
onnxruntime/test/contrib_ops/moe_packed_routing_test.cc |
Adds routing, malformed-route isolation, compatibility, and cancellation tests. |
onnxruntime/core/graph/contrib_ops/contrib_defs.cc |
Extends the schema with activation options and optional routing inputs. |
onnxruntime/contrib_ops/cuda/moe/moe_packed.h |
Adds activation and routing parameters. |
onnxruntime/contrib_ops/cuda/moe/moe_packed.cu |
Implements new routing paths, ReLU-square activation, and BF16 IQ2 arithmetic. |
onnxruntime/contrib_ops/cuda/moe/moe_packed.cc |
Validates new input combinations and configures CUDA dispatch. |
docs/OperatorKernels.md |
Updates the supported kernel inputs and types. |
docs/ContribOperators.md |
Documents the expanded operator schema. |
Review effort: Balanced
Comment on lines
+1524
to
+1525
| MoEPackedReluIq2ContiguousKKernel<false, 8><<<fc2_grid, 256, 0, stream>>>( | ||
| fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out); |
Rui Ren (rui-ren)
force-pushed
the
rui-ren/bolt-packed-moe-on-33953-20261107
branch
2 times, most recently
from
November 9, 2026 18:34
7b2b13a to
e50e349
Compare
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
Contributor
There was a problem hiding this comment.
Reviewed this layer (stack #33967, layer 2) at e50e34972b.
Implementation. I traced the validation, workspace/chunk accounting and each new kernel, and ran the PR-head kernel code in a standalone H200 harness: random IQ2_XXS payloads against an independent dequantized reference (decode and prefill tiles), sigmoid top-k including ties and NaN/Inf up to 512 experts, and the preselected-route validator; compute-sanitizer (memcheck, initcheck, racecheck) was clean. I did not find an implementation defect. The remaining risk is in tests that cannot fail and in undocumented runtime-only constraints.
Should be fixed in this layer
MalformedRoutesPoisonOnlyTheirTokendoes not verify NaN poisoning in BF16:TensorCheck<BFloat16>checksisnan(expected)instead of the actual value (inline on the test).- The "FP32 weights" fixtures give identical BF16 expectations if the weights are rounded to BF16 (inline), so they cannot catch narrowing.
- The two coverage requests from the earlier review are still open (asymmetric/independent-reference fixtures, and an FC2-specific cancellation fixture with
inter_size >= 1024); the checklist item in the PR description is unchecked. I added details to the first thread. - Please clarify the
compute_75overlay result. The PR description and the author review say a compute_75-only PTX overlay passedReluIq2PreservesContiguousKCancellation. With compute_75 PTX the MMA body is compiled out (MmaBf16is guarded by__CUDA_ARCH__ >= 800),GetMoEPackedBf16KernelSupportreports false (ptxVersion < 80), and dispatch falls back to the SIMT kernels. Running the PR-head kernels on an H200, this fixture yields 1024 on the Tensor-Core path but 256 on the SIMT path for rows 1/8/17/129 (the SIMT FC1 lane tree adds the unit term at k=128 to2^25at k=0 before the cancellation, so FC1 is 1 instead of 2). So the expected value is specific to the Tensor-Core dispatch, the same situation asRoundedBf16ContiguousKCancellationin #33962, andHasCudaEnvironment(800)does not guarantee that path. Which dispatch did the overlay exercise?
Other requests (inline): negative tests for the ~10 new validation rules, naming (relu_squared vs the existing relu2), undocumented runtime-only constraints, the defaults-unchanged test not reaching prefill tiles, and unreachable relu_squared branches in the generic Tensor-Core SwiGLU kernels.
Minor / notes
- No test executes the >2048-row chunk loop (second-chunk offsets for preselected
indices/weightsand sigmoidrouter_input, and workspace reuse across chunks). I read the offset math as correct; a 2049-row case per routing mode (k=2 keeps it small) would pin it. - FP16
relu_squarednow rounds decoded weights to FP16 (p.relu_squared || is_bf16), whereas FP16 SwiGLU keeps FP32-decoded weights. That is a reasonable choice but is not mentioned in the PR text or the docs. - Error texts: with the default softmax routing and
router_input/router_weight/router_biassupplied, the check reports "sigmoid routing requires ...", which points at the wrong fix; and the constructor text "internally routed mode requires k <= 16" is only true for softmax (sigmoid routing is also internal but allows k <= 32). The// 1. Softmax + TopKcomment inmoe_packed.cuis stale now that it also covers sigmoid and preselected routing. docs/PackedWeightFormats.md(not part of this diff) still says MoEPacked always has three weight tensors with independently chosen formats. Withrelu_squaredthere is nouptensor and gate/down must both be IQ2_XXS, and the weight-rounding contract of the new BF16 IQ2 path (weights rounded to the activation dtype, FP32 accumulation with contiguous-K partitions, BF16 rounding of FC1/ReLU-square/FC2 outputs) is undocumented, unlike the MatMulPacked rounding table.
3 of 4 tasks
Rui Ren (rui-ren)
force-pushed
the
rui-ren/bolt-packed-moe-on-33953-20261107
branch
from
November 9, 2026 23:43
e50e349 to
79067e3
Compare
3 of 4 tasks
Preserve default SwiGLU and internal routing. Add IQ2_XXS ReLU-square gate/down experts, unchanged preselected FP32 weights and expert IDs through k=32, and safe affected-token poisoning for malformed routes. Exercise analytic outputs, routing boundaries and existing defaults. Production-flag translation units compile. Runtime numerical, generation and throughput qualification remain pending. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Add an explicit projected sigmoid routing mode with FP32 selection-only correction bias and unbiased normalized routing weights. Preserve the existing softmax/SwiGLU defaults and preselected routing. Retain the source reduction order, tie-breaking and route scaling, with FP16 and BF16 regression fixtures. Validation: all five focused MoE CUDA tests pass on H100; all eight combined projection/MoE tests pass compute-sanitizer memcheck with zero skips or errors. Source hashes match the tested v3 build. This is operator-level compatibility-overlay evidence; full-model strict bitwise parity currently fails and speculative/TPS/release acceptance remain outstanding. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Use source-compatible FC1 and FC2 K partitions, MMA ordering, and BF16 store boundaries for the IQ2 ReLU-square path. Reuse existing expert routing and finalization, retain other dispatch paths, and exercise cancellation plus 32-route prefill subtiles. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Keep unit contributions outside large-magnitude tensor-core partials so the exact expected result does not depend on Ampere versus Hopper internal MMA accumulation. Preserve the extreme cancellation input, kernel arithmetic and zero tolerance. The updated current packed test overlay executes all ten projection and MoE regressions on H100 without skips. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Use the canonical CPU/CUDA/DirectML/WebGPU CI artifact for the MoE layer so schema and kernel documentation match the actual registrations. Preserve artifact provenance outside the repository. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: rui-ren <ruiren1225@gmail.com>
Keep upstream compiled-BF16 fallback, bounded launches, and checked workspace accounting while retaining contiguous ReLU-square reductions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Signed-off-by: rui-ren <ruiren1225@gmail.com>
Retain the upstream 16-expert softmax limit while allowing the existing preselected and projected-sigmoid paths to select up to 32 experts. Apply the same range to checked workspace sizing. Existing k=32 routing regressions expose the inherited guard mismatch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Signed-off-by: rui-ren <ruiren1225@gmail.com>
Address all eight PR33963 coworker review threads. Use independent CPU IQ2_XXS dequantization and double-precision reference matmuls with asymmetric FC1/FC2 dimensions and nonuniform packed payloads, activations and expert routes. Bound BF16 deviations to two ULPs and at most 0.1% differing elements (one for tiny outputs), without relaxing existing exact comparisons. Verify actual poisoned BF16 outputs and unaffected tokens, use routing fixtures that demonstrably distinguish FP32 from BF16 weights, add 22 negative contract cases, and cover default prefill plus SwiGLU preselected k=2/32. Canonicalize relu2 while retaining relu_squared, reject attribute-only incompatibilities at construction, remove unreachable generic Tensor-Core relu branches while retaining SIMT support, and document routing/format constraints in the schema and generated operator docs. Review discussions: 4517915471, 4517913786, 4517913804, 4517913813, 4517913820, 4517913826, 4517913833, 4517913839. Validation: isolated host compilation of moe_packed.cc, moe_packed_routing_test.cc and contrib_defs.cc using preserved focused overlay recipes; configured CLANGFORMAT lint and git diff --check pass. CUDA execution and sanitizer remain for the parent's coordinated combined rebuild; no GPU or shared build artifacts were used. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Signed-off-by: rui-ren <ruiren1225@gmail.com>
The all-invalid routing case poisons every token, so a verifier that assumes token 1 is finite rejects correct NaN output. Follow the expected NaN pattern across the entire tensor, explicitly checking actual NaNs and retaining exact, finite-value comparisons for valid outputs. No tolerance changes. Validation: complete MoE host test compilation, eight extracted CPU verifier acceptance/rejection cases, configured CLANGFORMAT and diff checks pass. GPU rerun remains with the parent integrated overlay. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: f4d0e1a3-a389-4cc8-9c85-0fd0e3ffce5b Signed-off-by: rui-ren <ruiren1225@gmail.com>
Rui Ren (rui-ren)
force-pushed
the
rui-ren/bolt-packed-moe-on-33953-20261107
branch
from
November 10, 2026 00:40
79067e3 to
97c77c1
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Extends
com.microsoft::MoEPacked(CUDA) for Nemotron Bolt's ReLU-square experts androuting contract without introducing a competing expert operator. Adds
activation_type="relu_squared", preselected expert IDs with FP32 mixture weights,and integrated projected-sigmoid routing. Correction bias affects selection only;
mixture weights come from unbiased sigmoid affinities and post-normalization scaling.
The existing softmax/SwiGLU defaults remain supported.
Preselected and projected-sigmoid routing support up to 32 selected experts; the
existing internal softmax path retains its 16-expert top-k limit. Validation keeps
routing modes mutually exclusive and isolates malformed routes to their token.
BF16 IQ2 FC1/FC2 preserve complete FP32 scale reconstruction, contiguous-K partitions,
BF16 storage boundaries, and FP32 routing/final reduction.
Dependency: #33962 in
native stack #33967;
the eventual landing target is
main. Packed weight layouts and GenAI source do not change.Summary of Changes
Operator Contract and Documentation
onnxruntime/core/graph/contrib_ops/contrib_defs.ccdocs/ContribOperators.md, docs/OperatorKernels.mdCUDA Implementation and Regression Tests
onnxruntime/contrib_ops/cuda/moe/moe_packed.{cc,h,cu}onnxruntime/test/contrib_ops/moe_packed_routing_test.ccTesting
Review-fix snapshot:
79067e34feddff33a1b49d1bb56cfd21d141004c.Added independent nonuniform CPU-reference coverage, true BF16 NaN verification, discriminating FP32 route weights, 22 invalid-configuration cases and missing SwiGLU paths. Removed dead TC ReLU branches.
relu2is canonical;relu_squaredremains a compatibility alias. Expanded the authoritative schema and generated reference.Owning-layer validation: 10 actual focused CUDA tests (including 22 negative configurations); zero failures/skips.
Integrated tip regression: 51 C++ provider/schema tests, 5 actual internal CUDA tests and 2 Python IOBinding tests passed, zero failures/skips. All owning-layer provider cases are retained. The final overlay recompiles the changed CUDA/host/schema/test sources over ABI-matched preserved ORT support; not a clean full-source build, full-CI approval, independent HF parity or new end-to-end qualification. Targeted configured clang-format/Ruff and whitespace checks passed.
Historical measurements and qualification evidence (not current-head acceptance)
These previously reported results are retained for provenance. They have not been rerun on the rebased head above; earlier CI/build status and benchmark results do not establish current-head CI, accuracy, or merge readiness.
Measured CUDA kernel latency - 2026-11-07
BF16, IQ2_XXS, ReLU-square experts, latent hidden dimension 1024, top-k 9. Internal softmax routing versus preselected FP32 indices/weights using identical routes and the same corrected expert arithmetic. Selected logits are +16 and others -16; selected weights are 1/9. Shared/distinct labels describe expert reuse across rows (the one-row variants are equivalent workload controls). Sixteen rows create 144 routes, crossing the 128-route prefill-tiles dispatch threshold.
The ratio measures routing-path overhead removed, not this entire PR versus its parent. All 24 complete BF16 outputs are bitwise equal in both runs. Benefits diminish for distinct-expert, bandwidth-heavy batches; all cases, including those smaller gains, are shown.
H100 80GB HBM3, CUDA 12.8.93, driver 580.159.04, GCC 11.4.0. CUDA-event median per dispatch; 64 dispatches per captured CUDA graph, 50 warmups, three graph warmups, and 15 paired timed batches alternating AB/BA. Four rotating allocation banks contain deterministic synthetic weights (seed 617); this is not a cold-cache guarantee or a checkpoint-weight benchmark. Raw p95 values describe batch-average latency, not individual-invocation tails. No profiler runs in the timing window. The assigned GPU had empty admission, one stable compute client during execution, and clean release afterward; clocks were not locked. Two independent 39-case runs passed, including the published replay runner.
Scope: same-build production CUDA dispatch alternatives, not historical parent-commit binaries. The measured runtime is an isolated CUDA provider relink with preserved core/support, not a clean full-source ORT build or end-to-end model qualification. Each measured operator's CUDA source hash matches its owning PR head. A clean source build is a separate replication variant and must be measured again. Kernel-only timings do not establish model TTFT/tok/s, MMLU/LCB, speculative integration, or full-model readiness.
Complete primary/repeat results, numerical metrics, source hashes and benchmark code. All 39 cases passed in both supervised runs; full-model work remains paused. No source push, stack change, merge, auto-merge or ready-for-review promotion is included.
Checklist